"""
@author: chenzhenhua
@project: jf_fashion
@file: pruning.py
@time: 2021/8/2 0002 17:11
@desc:
"""

activation = {}
def get_activation(name):
    def hook(model, input, output):
        # 如果你想feature的梯度能反向传播，那么去掉 detach（）
        activation[name] = output.detach()
    return hook